package mmult

import chisel3._
import chisel3.util._
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}

class MatrixMultiplication2(M: Int, N: Int, P: Int) extends Module {
  val io = IO(new MatrixMultiplication2IO(M, N, P))

  io.out := DontCare
  io.outValid := false.B

  private val sLoad :: sCompute :: sStore :: Nil = Enum(3)

  private val state = RegInit(sLoad)

  val regA = Reg(Vec(M * P, SInt(32.W)))
  val regB = Reg(Vec(P * N, SInt(32.W)))
  val regC = Reg(Vec(M * N, SInt(32.W)))

  val i = Counter(M)
  val j = Counter(N)

  switch(state) {
    is(sLoad) {
      when(io.inValid) {
        regA := io.a
        regB := io.b
        state := sCompute
      }
    }
    is(sCompute) {
      var sum = 0.S(32.W)

      for (k <- 0 until P) {
        sum = sum + regA(i.value * P.U + k.U) * regB(k.U * N.U + j.value)
      }

      regC(i.value * N.U + j.value) := sum

      when(j.inc() && i.inc()) {
        state := sStore
      }
    }
    is(sStore) {
      io.out := regC
      io.outValid := true.B
    }
  }
}

object MatrixMultiplication2 extends App {
  (new ChiselStage).execute(
    Array("-X", "verilog", "-td", "source/"),
    Seq(
      ChiselGeneratorAnnotation(() => new MatrixMultiplication2(3, 2, 5))
    )
  )
}
